package com.ma5d.es.integration;

import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.http.HttpHost;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;
import org.junit.Test;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;

public class ElasticsearchFlinkTest {

    @Test
    public void flinkIntegration() throws Exception {
        StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<String> dataStreamSource = executionEnvironment.socketTextStream("localhost", 9999);
        List<HttpHost> httpHostList = Collections.singletonList(new HttpHost("127.0.0.1", 9200, "http"));
        ElasticsearchSinkFunction<String> elasticsearchSinkFunction = (s, runtimeContext, requestIndexer) -> {
            HashMap<Object, Object> jsonMap = new HashMap<>();
            jsonMap.put("data", s);
            IndexRequest indexRequest = Requests.indexRequest();
            indexRequest.index("flink-index");
            indexRequest.id("9001");
            indexRequest.source(jsonMap);
        };
        ElasticsearchSink.Builder<String> elasticsearchSinkBuilder = new ElasticsearchSink.Builder<>(httpHostList, elasticsearchSinkFunction);
        elasticsearchSinkBuilder.setBulkFlushMaxActions(1);
        dataStreamSource.addSink(elasticsearchSinkBuilder.build());
        executionEnvironment.execute("flink-es");
    }
}
